import onnx
from nni.compression.onnx import apply_pruning, generate_masks
from nni.compression.onnx.utils import get_masks_code
from nni.compression.onnx.pruner import L1NormPruner
import numpy as np

# 加载预训练的ONNX模型
model = onnx.load('/disk2/xd/project/yolov5/runs/train/exp12/weights/best.onnx')

# 定义剪枝配置 - 对所有卷积层应用50%的剪枝率
config_list = [{
    'sparsity': 0.5,
    'op_types': ['Conv']
}]

# 创建L1范数剪枝器
pruner = L1NormPruner(model, config_list)

# 生成剪枝掩码
masks = generate_masks(pruner, model)

# 应用剪枝到模型
pruned_model = apply_pruning(model, masks)

# 保存剪枝后的ONNX模型
onnx.save(pruned_model, '/disk2/xd/project/yolov5/runs/train/exp12/weights/best_pruned.onnx')

# 打印剪枝统计信息
def print_pruning_statistics(original_model, pruned_model):
    original_nodes = len(original_model.graph.node)
    pruned_nodes = len(pruned_model.graph.node)
    original_params = sum(np.prod(np.array(param.dims)) for param in original_model.graph.initializer)
    pruned_params = sum(np.prod(np.array(param.dims)) for param in pruned_model.graph.initializer)
    
    print(f"原始模型节点数: {original_nodes}")
    print(f"剪枝后模型节点数: {pruned_nodes}")
    print(f"参数减少比例: {(original_params - pruned_params) / original_params * 100:.2f}%")
    print(f"节点减少比例: {(original_nodes - pruned_nodes) / original_nodes * 100:.2f}%")

print_pruning_statistics(model, pruned_model)

# 生成剪枝掩码的Python代码（可选）
# 这可以帮助你了解具体的剪枝策略
masks_code = get_masks_code(masks)
with open('pruning_masks.py', 'w') as f:
    f.write(masks_code)